package fr.inserm.umr915.vcfannotator.annotators;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.URL;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.regex.Pattern;
import java.util.zip.GZIPInputStream;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import fr.inserm.umr915.libjson.ArrayNode;
import fr.inserm.umr915.libjson.NodeFactory;
import fr.inserm.umr915.libjson.ObjectNode;
import fr.inserm.umr915.vcfannotator.AnnotationException;
import fr.inserm.umr915.vcfannotator.Variant;
import fr.inserm.umr915.vcfannotator.bio.GeneticCode;
import fr.inserm.umr915.vcfannotator.bio.GenomicSequence;
import fr.inserm.umr915.vcfannotator.bio.KnownGene;
import fr.inserm.umr915.vcfannotator.bio.MutatedSequence;
import fr.inserm.umr915.vcfannotator.bio.ProteinSequence;
import fr.inserm.umr915.vcfannotator.bio.SequenceProvider;


import com.sleepycat.je.DatabaseEntry;
import com.sleepycat.je.LockMode;
import com.sleepycat.je.OperationStatus;
import com.sleepycat.je.SecondaryCursor;

public class EffectPredictorAnnotator
	extends AbstractAnnotator
	{
	static Log LOG=LogFactory.getLog("vcf.annotator");
	
	
	static final String KEY_TYPE="type";
	static final String KEY_SPLICING="splicing";

	private Map<String, List<KnownGene>> chrom2genes=new HashMap<String, List<KnownGene>>();
	boolean loadWholeSegment=false;
	private SequenceProvider sequenceProvider;
	private String build=null;
	private boolean loadWholeChromosome=true;
	private Map<String, GeneticCode> geneticCodeMap=new HashMap<String, GeneticCode>();
	
	public EffectPredictorAnnotator()
		{		
		
		}
	
	public String getBuild()
		{
		return build;
		}
	
	public void setBuild(String build)
		{
		this.build = build;
		}
	
	public SequenceProvider getSequenceProvider()
		{
		return sequenceProvider;
		}
	
	public void setSequenceProvider(SequenceProvider sequenceProvider)
		{
		this.sequenceProvider = sequenceProvider;
		}
	
	public void setLoadingWholeChromosome(boolean loadWholeChromosome)
		{
		this.loadWholeChromosome = loadWholeChromosome;
		}
	
	public boolean isLoadingWholeChromosome()
		{
		return loadWholeChromosome;
		}
	
	
	
	@Override
	public void annotate() throws AnnotationException
		{
		NodeFactory nodeFactory=new NodeFactory();
		Variant.Binding varBinding=new Variant.Binding();
		SecondaryCursor cursor=null;
		try
			{
			loadKnownGenes();
			cursor=super.baseDatabase.getVariantDatabase().openCursor(null, null);
			final int extra=10000;
			int nSeen=0;
			GenomicSequence genomicSeq=null;
			DatabaseEntry varKey=new DatabaseEntry();
			DatabaseEntry pKey=new DatabaseEntry();
			DatabaseEntry data=new DatabaseEntry();
			
			
			
		    while(cursor.getNext(varKey, pKey,data, LockMode.DEFAULT)==OperationStatus.SUCCESS)
	            {
		    	Variant call=varBinding.entryToObject(varKey);
		    	if(!call.getBuild().equals(getBuild())) continue;
		    	
		    	List<KnownGene> genes= getGenes(call);
	            
	            LOG.info(genes.size());
	            
	            if(genes.isEmpty())
	            	{
	            	LOG.info("GENOMIC");
	            	continue;
	            	}
		    	
		    	ArrayNode predictionsNode=nodeFactory.createArrayNode();
	            LOG.info(call.toString());
	            int position= call.getPosition()-1;
	            String ref=call.getRef().toUpperCase();
	        	String alts=call.getAlt().toUpperCase();
	        	
	        	
	        	
	        	
	        	for(String alt: Variant.alts(ref, alts))
		        	{
		    		LOG.info("var:"+ref+" "+alt);
		            
		            for(KnownGene gene:genes)
		            	{
		            	ObjectNode kgNode= nodeFactory.createObjectNode();
		            	kgNode.put("alt", alt);
		            	ArrayNode typeNode=nodeFactory.createArrayNode();
		            	kgNode.put("kgId", gene.getName());
		            	kgNode.put("gene", gene.getGeneSymbol());
		            	kgNode.put("types", typeNode);
		            	LOG.info(gene.getName());
		            	predictionsNode.add(kgNode);
		            	
		            	//switch to 0 based coordinate
		        		
		        		
		        		
		            	if( (ref.equals("A") || ref.equals("T") || ref.equals("G") || ref.equals("C")) &&
		            		(alt.equals("A") || alt.equals("T") || alt.equals("G") || alt.equals("C"))
		            		)
			        		{
			        		LOG.info("fetch genome ");
		            		GeneticCode geneticCode= getGeneticCodeByChromosome(gene.getChromosome());
			        		StringBuilder wildRNA=null;
			        		ProteinSequence wildProt=null;
			        		ProteinSequence mutProt=null;
			        		MutatedSequence mutRNA=null;
			        		int position_in_cdna=-1;
			        		
			        		
			        		if(genomicSeq==null ||
			        	               !gene.getChromosome().equals(genomicSeq.getChrom()) ||
			        	               !(genomicSeq.getChromStart()<=gene.getTxStart() && gene.getTxEnd()<= genomicSeq.getChromEnd())
			        	               )
		    	            	{
			        			if(isLoadingWholeChromosome())
			        				{
			        				genomicSeq=this.sequenceProvider.getGenomicSequence(
				    	            		gene.getChromosome()
				    	            		);
			        				}
			        			else
			        				{
			        				genomicSeq=this.sequenceProvider.getGenomicSequence(
			    	            		gene.getChromosome(),
			    	            		Math.max(gene.getTxStart()-extra,0),
			    	            		gene.getTxEnd()+extra
			    	            		);
			        				}
		    	            	}
			        		
			        		if(!String.valueOf(genomicSeq.charAt(position)).equalsIgnoreCase(ref))
			        			{
			        			System.err.println("Warning REF!=GENOMIC SEQ!!! at "+genomicSeq.charAt(position)+"/"+ref);
			        			return;
			        			}
			        		
			        		if(gene.isForward())
			            		{
			            		if(position < gene.getCdsStart())
			            			{
			            			typeNode.insert("UTR5");
			            			continue;
			            			}
			            		else if( gene.getCdsEnd()<= position )
			            			{
			            			typeNode.insert("UTR3");
			            			continue;
			            			}
			            		
			            		int exon_index=0;
			            		while(exon_index< gene.getExonCount())
			            			{
			            			KnownGene.Exon exon= gene.getExon(exon_index);
			            			for(int i= exon.getStart();
			            					i< exon.getEnd();
			            					++i)
			            				{
			            				if(i==position)
			        						{
			        						kgNode.put("exon.name", exon.getName());
			        						}
			            				if(i< gene.getCdsStart()) continue;
			            				if(i>=gene.getCdsEnd()) break;
			        					
			        					if(wildRNA==null)
			        						{
			        						wildRNA=new StringBuilder();
			        						mutRNA=new MutatedSequence(wildRNA);
			        						}
			        					
			        					if(i==position)
			        						{
			        						typeNode.insert("EXON");
			        						kgNode.put("exon.name",exon.getName());
			        						position_in_cdna=wildRNA.length();
			        						kgNode.put("pos.cdna", String.valueOf(position_in_cdna));
			        						//in splicing ?
			        						if(exon.isSplicing(position))
			        							{
			        							typeNode.insert("SPLICING");
			        							
			        							if(exon.isSplicingAcceptor(position))
			        								{
			        								typeNode.insert("SPLICING_ACCEPTOR");
			        								}
			        							else  if(exon.isSplicingDonor(position))
			        								{
			        								typeNode.insert("SPLICING_DONOR");
			        								}
			        							}
			        						}
			        					
			            				wildRNA.append(genomicSeq.charAt(i));
			            				
			            				if(i==position)
			            					{
			            					mutRNA.put(
			            							position_in_cdna,
			            							alt.charAt(0)
			            							);
			            					}
			            				
			            				if(wildRNA.length()%3==0 && wildRNA.length()>0 && wildProt==null)
				            				{
				            				wildProt=new ProteinSequence(geneticCode,wildRNA);
				            				mutProt=new ProteinSequence(geneticCode,mutRNA);
				            				}
			            				}
			            			KnownGene.Intron intron= exon.getNextIntron();
			            			if(intron!=null && intron.contains(position))
			            				{
			            				kgNode.put("intron.name",intron.getName());
			            				typeNode.insert("INTRON");
			            				
			            				if(intron.isSplicing(position))
			        						{
			            					kgNode.put(KEY_SPLICING, "INTRON_SPLICING");
			        						if(intron.isSplicingAcceptor(position))
			        							{
			        							typeNode.insert("INTRON_SPLICING_ACCEPTOR");
			        							}
			        						else if(intron.isSplicingDonor(position))
			        							{
			        							typeNode.insert("INTRON_SPLICING_DONOR");
			        							}
			        						}
			            				}
			            			++exon_index;
			            			}
			            		
			            		
			            		
			            		}
			            	else // reverse orientation
			            		{
			            	
			            		if(position < gene.getCdsStart())
			            			{
			            			typeNode.insert("UTR3");
			            			continue;
			            			}
			            		else if( gene.getCdsEnd()<=position )
			            			{
			            			typeNode.insert("UTR5");
			            			continue;
			            			}
			            	
			            		
			            		int exon_index = gene.getExonCount()-1;
			            		while(exon_index >=0)
			            			{
			            			KnownGene.Exon exon= gene.getExon(exon_index);
			            			for(int i= exon.getEnd()-1;
			            				    i>= exon.getStart();
			            				--i)
			            				{
			            				if(i==position)
			        						{
			            					kgNode.put("exon.name", exon.getName());
			        						}
			            				if(i>= gene.getCdsEnd()) continue;
			            				if(i<  gene.getCdsStart()) break;
			            				
			            				if(wildRNA==null)
			        						{
			        						wildRNA=new StringBuilder();
			        						mutRNA=new MutatedSequence(wildRNA);
			        						}
			            				
			            				if(i==position)
			        						{
			            					typeNode.insert("EXON");
			            					position_in_cdna=wildRNA.length();
			        						kgNode.put("pos.cdna",String.valueOf(position_in_cdna));
			        						//in splicing ?
			        						if(exon.isSplicing(position))
			        							{
			        							typeNode.insert("INTRON_SPLICING");
			        							
			        							if(exon.isSplicingAcceptor(position))
			        								{
			        								typeNode.insert("INTRON_SPLICING_ACCEPTOR");
			        								}
			        							else  if(exon.isSplicingDonor(position))
			        								{
			        								typeNode.insert("INTRON_SPLICING_DONOR");
			        								}
			        							}
			        						
			        						
			        						mutRNA.put(
			        								position_in_cdna,
			        								complement(alt.charAt(0))
			        								);
			        						}
			            				
			            				wildRNA.append(complement(genomicSeq.charAt(i)));
			            				if( wildRNA.length()%3==0 &&
			            					wildRNA.length()>0 &&
			            					wildProt==null)
				            				{
				            				wildProt=new ProteinSequence(geneticCode,wildRNA);
				            				mutProt=new ProteinSequence(geneticCode,mutRNA);
				            				}
			            				
			            				}
			            			
			            			KnownGene.Intron intron= exon.getPrevIntron();
			            			if(intron!=null &&
			            				intron.contains(position))
			            				{
			            				kgNode.put("intron.name",intron.getName());
			            				typeNode.insert("INTRON");
			            				
			            				if(intron.isSplicing(position))
			        						{
			            					typeNode.insert("INTRON_SPLICING");
			        						if(intron.isSplicingAcceptor(position))
			        							{
			        							typeNode.insert("INTRON_SPLICING_ACCEPTOR");
			        							}
			        						else if(intron.isSplicingDonor(position))
			        							{
			        							typeNode.insert("INTRON_SPLICING_DONOR");
			        							}
			        						}
			            				}
			            			--exon_index;
			            			}
		
			            		}//end of if reverse
			        		if( wildProt!=null &&
			        			mutProt!=null && 
			        			position_in_cdna>=0)
				    			{
			            		int pos_aa=position_in_cdna/3;
			            		int mod= position_in_cdna%3;
			            		kgNode.put("wild.codon",""+
			            			wildRNA.charAt(position_in_cdna-mod+0)+
			            			wildRNA.charAt(position_in_cdna-mod+1)+
			            			wildRNA.charAt(position_in_cdna-mod+2)
			            			);
			            		kgNode.put("mut.codon",""+
			            			mutRNA.charAt(position_in_cdna-mod+0)+
			            			mutRNA.charAt(position_in_cdna-mod+1)+
			            			mutRNA.charAt(position_in_cdna-mod+2)
			            			);
			            		kgNode.put("position.protein",String.valueOf(pos_aa+1));
			            		kgNode.put("wild.aa",String.valueOf(wildProt.charAt(pos_aa)));
			            		kgNode.put("mut.aa",String.valueOf(mutProt.charAt(pos_aa)));
				    			if(isStop(wildProt.charAt(pos_aa)) &&
				    			   !isStop(mutProt.charAt(pos_aa)))
				    				{
				    				typeNode.insert("EXON_STOP_LOST");
				    				}
				    			else if( !isStop(wildProt.charAt(pos_aa)) &&
				    				 isStop(mutProt.charAt(pos_aa)))
				    				{
				    				typeNode.insert("EXON_STOP_GAINED");
				    				}
				    			else if(wildProt.charAt(pos_aa)==mutProt.charAt(pos_aa))
				    				{
				    				typeNode.insert("EXON_CODING_SYNONYMOUS");
				    				}
				    			else
				    				{
				    				typeNode.insert("EXON_CODING_NON_SYNONYMOUS");
				    				}
				    			}
			        		}//end of simpe ATCG
		            	else
		            		{
			        		Integer wildrna=null;
			        		int position_in_cdna=-1;
			        		
			        		
			        		
			        		if(gene.isForward())
			            		{
			            		if(position < gene.getCdsStart())
			            			{
			            			typeNode.insert("UTR5");
			            			continue;
			            			}
			            		else if( gene.getCdsEnd()<= position )
			            			{
			            			typeNode.insert("UTR3");
			            			continue;
			            			}
			            		
			            		int exon_index=0;
			            		while(exon_index< gene.getExonCount())
			            			{
			            			KnownGene.Exon exon= gene.getExon(exon_index);
			            			for(int i= exon.getStart();
			            					i< exon.getEnd();
			            					++i)
			            				{
			            				if(i==position)
			        						{
			        						kgNode.put("exon.name", exon.getName());
			        						}
			            				if(i< gene.getCdsStart()) continue;
			            				if(i>=gene.getCdsEnd()) break;
			        					
			        					if(wildrna==null)
			        						{
			        						wildrna=0;
			        						}
			        					
			        					if(i==position)
			        						{
			        						typeNode.insert("EXON");
			        						kgNode.put("exon.name",exon.getName());
			        						position_in_cdna=wildrna;
			        						kgNode.put("pos.cdna", String.valueOf(wildrna));
			        						//in splicing ?
			        						if(exon.isSplicing(position))
			        							{
			        							typeNode.insert("SPLICING");
			        							
			        							if(exon.isSplicingAcceptor(position))
			        								{
			        								typeNode.insert("SPLICING_ACCEPTOR");
			        								}
			        							else  if(exon.isSplicingDonor(position))
			        								{
			        								typeNode.insert("SPLICING_DONOR");
			        								}
			        							}
			        						}
			        					
			        					wildrna++;
			            				}
			            			KnownGene.Intron intron= exon.getNextIntron();
			            			if(intron!=null && intron.contains(position))
			            				{
			            				kgNode.put("intron.name",intron.getName());
			            				typeNode.insert("INTRON");
			            				
			            				if(intron.isSplicing(position))
			        						{
			            					typeNode.insert("INTRON_SPLICING");
			        						if(intron.isSplicingAcceptor(position))
			        							{
			        							typeNode.insert("INTRON_SPLICING_ACCEPTOR");
			        							}
			        						else if(intron.isSplicingDonor(position))
			        							{
			        							typeNode.insert("INTRON_SPLICING_DONOR");
			        							}
			        						}
			            				}
			            			++exon_index;
			            			}
			            		}
			            	else // reverse orientation
			            		{
			            	
			            		if(position < gene.getCdsStart())
			            			{
			            			typeNode.insert("UTR3");
			            			continue;
			            			}
			            		else if( gene.getCdsEnd()<=position )
			            			{
			            			typeNode.insert("UTR5");
			            			continue;
			            			}
			            	
			            		
			            		int exon_index = gene.getExonCount()-1;
			            		while(exon_index >=0)
			            			{
			            			KnownGene.Exon exon= gene.getExon(exon_index);
			            			for(int i= exon.getEnd()-1;
			            				    i>= exon.getStart();
			            				--i)
			            				{
			            				if(i==position)
			        						{
			            					kgNode.put("exon.name", exon.getName());
			        						}
			            				if(i>= gene.getCdsEnd()) continue;
			            				if(i<  gene.getCdsStart()) break;
			            				
			            				if(wildrna==null)
			        						{
			        						wildrna=0;
			        						}
			            				
			            				if(i==position)
			        						{
			            					typeNode.insert("EXON");
			            					position_in_cdna=wildrna;
			        						kgNode.put("pos.cdna",String.valueOf(position_in_cdna));
			        						//in splicing ?
			        						if(exon.isSplicing(position))
			        							{
			        							typeNode.insert("INTRON_SPLICING");
			        							
			        							if(exon.isSplicingAcceptor(position))
			        								{
			        								typeNode.insert("INTRON_SPLICING_ACCEPTOR");
			        								}
			        							else  if(exon.isSplicingDonor(position))
			        								{
			        								typeNode.insert("INTRON_SPLICING_DONOR");
			        								}
			        							}
			        						}
			            				
			            				wildrna++;
			            				}
			            			
			            			KnownGene.Intron intron= exon.getPrevIntron();
			            			if(intron!=null &&
			            				intron.contains(position))
			            				{
			            				kgNode.put("intron.name",intron.getName());
			            				typeNode.insert("INTRON");
			            				
			            				if(intron.isSplicing(position))
			        						{
			            					typeNode.insert("INTRON_SPLICING");
			        						if(intron.isSplicingAcceptor(position))
			        							{
			        							typeNode.insert("INTRON_SPLICING_ACCEPTOR");
			        							}
			        						else if(intron.isSplicingDonor(position))
			        							{
			        							typeNode.insert("INTRON_SPLICING_DONOR");
			        							}
			        						}
			            				}
			            			--exon_index;
			            			}
		
			            		}//end of if reverse
			        		if( wildrna!=null &&
			        			position_in_cdna>=0)
				    			{
			            		int pos_aa=position_in_cdna/3;
			            		kgNode.put("position.protein",String.valueOf(pos_aa+1));
				    			}
		            		}
		            	}
		            
		            
		            if(predictionsNode!=null && predictionsNode.size()!=0)
		            	{
		            	LOG.info(predictionsNode);
		    			ObjectNode annotation= nodeBinding.entryToObject(data).asObjectNode();
		    			annotation.put("effect", predictionsNode);
		    			
		    			nodeBinding.objectToEntry(annotation, data);
		    			if(baseDatabase.getJsonDatabase().put(null, pKey, data)!=OperationStatus.SUCCESS)
		    				{
		    				throw new RuntimeException("Cannot re-insert "+call+":"+annotation);
		    				}
		            	}
		        	}
	           }
		   
			
			
			}catch(IOException err)
			{
			throw new AnnotationException(err);
			}
		finally
			{
			if(cursor!=null) cursor.close();
			}
		}
	@Override
	public void cleanup()
		{
		this.chrom2genes.clear();
		if(this.sequenceProvider!=null) this.sequenceProvider.close();
		}
	
	
	
	

	

	
	private static char complement(char c)
		{
		switch(c)
			{
			case 'A': return 'T';
			case 'T': return 'A';
			case 'G': return 'C';
			case 'C': return 'G';
			default:throw new IllegalArgumentException(""+c);
			}
		}
	
	private BufferedReader openReader(String uri)  throws IOException
		{
		URL url=new URL(uri);
		return new BufferedReader(new InputStreamReader(new GZIPInputStream(url.openStream())));
		}
	
	public void loadKnownGenes() throws IOException
		{
		Map<String, KnownGene> kgId2gene=new HashMap<String, KnownGene>();
		Pattern tab=Pattern.compile("\t");
		Set<String> chromosomes= super.baseDatabase.getChromosomes(getBuild());
		BufferedReader in=openReader("http://hgdownload.cse.ucsc.edu/goldenPath/"+getBuild()+"/database/knownGene.txt.gz");
		String line;
		while((line=in.readLine())!=null)
			{
			String tokens[]=tab.split(line);
			if(!chromosomes.contains(tokens[1])) continue;
			KnownGene g=new KnownGene(tokens);
			List<KnownGene> L=this.chrom2genes.get(g.getChromosome());
			if(L==null)
				{
				L=new ArrayList<KnownGene>();
				this.chrom2genes.put(g.getChromosome(),L);
				}
			L.add(g);
			kgId2gene.put(g.getName(),g);
			}
		in.close();
		
		for(String chr:this.chrom2genes.keySet())
			{
			List<KnownGene> L=this.chrom2genes.get(chr);
			LOG.info(chr+":"+L.size());
			Collections.sort(L,
				new java.util.Comparator<KnownGene>()
					{
					@Override
					public int compare(KnownGene o1, KnownGene o2)
						{
						return o1.getTxStart()-o2.getTxStart();
						}
					});
			}
		

		in=openReader("http://hgdownload.cse.ucsc.edu/goldenPath/"+getBuild()+"/database/kgXref.txt.gz");
		while((line=in.readLine())!=null)
			{
			String tokens[]=tab.split(line);
			KnownGene g=kgId2gene.get(tokens[0]);
			if(g==null) continue;
			g.setGeneSymbol(tokens[4]);
			}
		in.close();
		
		}
	
	
	private List<KnownGene> getGenes(Variant variant)
		{
		List<KnownGene> array = new ArrayList<KnownGene>();
		List<KnownGene> genes = this.chrom2genes.get(variant.getChrom());
		if(genes==null)
			{
			LOG.info("no genes for "+variant.getChrom());
			return array;
			}
		
		int first=0;
	   
	    while(first< genes.size())
	    	{
	    	KnownGene gene= genes.get(first);
	    	if( gene.getTxStart()<= (variant.getPosition()-1 ) &&
	    		(variant.getPosition()-1)< gene.getTxEnd())
	    		{
	    		array.add(gene);
	    		}
	    	++first;
	    	}
		return array;
		}
	
	public void setGeneticCodeMap(Map<String, GeneticCode> geneticCodeMap)
		{
		this.geneticCodeMap = geneticCodeMap;
		}
	
	
	public Map<String, GeneticCode> getGeneticCodeMap()
		{
		return geneticCodeMap;
		}
	
	public GeneticCode getGeneticCodeByChromosome(String chrom)
		{
		GeneticCode c= getGeneticCodeMap().get(chrom);
		if(c==null) return new GeneticCode();
		return c;
		}	
	
	private boolean isStop(char aa)
		{
		return aa=='*';
		}
	}
